{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/tomasonjo/blogs/blob/master/llm/graph_semantic_layer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-yk_vT3wH2zQ",
    "outputId": "bdd88072-5239-44c0-90ca-71fe4b996bfd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: langchain in /usr/local/lib/python3.10/dist-packages (0.0.354)\n",
      "Requirement already satisfied: openai in /usr/local/lib/python3.10/dist-packages (1.6.1)\n",
      "Requirement already satisfied: neo4j in /usr/local/lib/python3.10/dist-packages (5.16.0)\n",
      "Requirement already satisfied: PyYAML>=5.3 in /usr/local/lib/python3.10/dist-packages (from langchain) (6.0.1)\n",
      "Requirement already satisfied: SQLAlchemy<3,>=1.4 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.0.23)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /usr/local/lib/python3.10/dist-packages (from langchain) (3.9.1)\n",
      "Requirement already satisfied: async-timeout<5.0.0,>=4.0.0 in /usr/local/lib/python3.10/dist-packages (from langchain) (4.0.3)\n",
      "Requirement already satisfied: dataclasses-json<0.7,>=0.5.7 in /usr/local/lib/python3.10/dist-packages (from langchain) (0.6.3)\n",
      "Requirement already satisfied: jsonpatch<2.0,>=1.33 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.33)\n",
      "Requirement already satisfied: langchain-community<0.1,>=0.0.8 in /usr/local/lib/python3.10/dist-packages (from langchain) (0.0.8)\n",
      "Requirement already satisfied: langchain-core<0.2,>=0.1.5 in /usr/local/lib/python3.10/dist-packages (from langchain) (0.1.6)\n",
      "Requirement already satisfied: langsmith<0.1.0,>=0.0.77 in /usr/local/lib/python3.10/dist-packages (from langchain) (0.0.77)\n",
      "Requirement already satisfied: numpy<2,>=1 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.23.5)\n",
      "Requirement already satisfied: pydantic<3,>=1 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.10.13)\n",
      "Requirement already satisfied: requests<3,>=2 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.31.0)\n",
      "Requirement already satisfied: tenacity<9.0.0,>=8.1.0 in /usr/local/lib/python3.10/dist-packages (from langchain) (8.2.3)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in /usr/local/lib/python3.10/dist-packages (from openai) (3.7.1)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /usr/lib/python3/dist-packages (from openai) (1.7.0)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from openai) (0.26.0)\n",
      "Requirement already satisfied: sniffio in /usr/local/lib/python3.10/dist-packages (from openai) (1.3.0)\n",
      "Requirement already satisfied: tqdm>4 in /usr/local/lib/python3.10/dist-packages (from openai) (4.66.1)\n",
      "Requirement already satisfied: typing-extensions<5,>=4.7 in /usr/local/lib/python3.10/dist-packages (from openai) (4.9.0)\n",
      "Requirement already satisfied: pytz in /usr/local/lib/python3.10/dist-packages (from neo4j) (2023.3.post1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (23.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (6.0.4)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.9.4)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.4.1)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.3.1)\n",
      "Requirement already satisfied: idna>=2.8 in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.5.0->openai) (3.6)\n",
      "Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.5.0->openai) (1.2.0)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json<0.7,>=0.5.7->langchain) (3.20.1)\n",
      "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json<0.7,>=0.5.7->langchain) (0.9.0)\n",
      "Requirement already satisfied: certifi in /usr/local/lib/python3.10/dist-packages (from httpx<1,>=0.23.0->openai) (2023.11.17)\n",
      "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.10/dist-packages (from httpx<1,>=0.23.0->openai) (1.0.2)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /usr/local/lib/python3.10/dist-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai) (0.14.0)\n",
      "Requirement already satisfied: jsonpointer>=1.9 in /usr/local/lib/python3.10/dist-packages (from jsonpatch<2.0,>=1.33->langchain) (2.4)\n",
      "Requirement already satisfied: packaging<24.0,>=23.2 in /usr/local/lib/python3.10/dist-packages (from langchain-core<0.2,>=0.1.5->langchain) (23.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (3.3.2)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (2.0.7)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from SQLAlchemy<3,>=1.4->langchain) (3.0.2)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7,>=0.5.7->langchain) (1.0.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install langchain openai neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "wxfYE2M5H8BE"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from langchain.graphs import Neo4jGraph\n",
    "# credentials\n",
    "url = \"\"\n",
    "username = \"neo4j\"\n",
    "password = \"\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-\"\n",
    "\n",
    "graph = Neo4jGraph(url=url, username=username, password=password)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "U7AuMk_iPAEG",
    "outputId": "c5f4d733-08ab-4bcc-948a-0cc84df9d9bb"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.query(\n",
    "    \"CREATE FULLTEXT INDEX movie IF NOT EXISTS FOR (m:Movie) ON EACH [m.title]\"\n",
    ")\n",
    "graph.query(\n",
    "    \"CREATE FULLTEXT INDEX person IF NOT EXISTS FOR (p:Person) ON EACH [p.name]\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-zv8U_qWPWoT",
    "outputId": "713fb4fa-81a0-4025-ead3-b1ad14d7328d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pulp~0.8 AND Fiction~0.8\n"
     ]
    }
   ],
   "source": [
    "def generate_full_text_query(input: str) -> str:\n",
    "    full_text_query = \"\"\n",
    "    words = input.split()\n",
    "    for word in words[:-1]:\n",
    "        full_text_query += f\" {word}~0.8 AND\"\n",
    "    full_text_query += f\" {words[-1]}~0.8\"\n",
    "    return full_text_query.strip()\n",
    "\n",
    "print(generate_full_text_query(\"Pulp Fiction\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yEQbjwPvA4C3",
    "outputId": "4dbf933d-3c2f-44a2-e68b-2e02cbfe13a3"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'candidate': 'Matrix, The', 'label': 'Movie'},\n",
       " {'candidate': 'Matrix Revolutions, The', 'label': 'Movie'},\n",
       " {'candidate': 'Matrix Reloaded, The', 'label': 'Movie'}]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import List, Dict, Optional\n",
    "\n",
    "candidate_query = \"\"\"\n",
    "CALL db.index.fulltext.queryNodes($index, $fulltextQuery)\n",
    "YIELD node\n",
    "RETURN coalesce(node.name, node.title) AS candidate, [el in labels(node) WHERE el IN ['Person', 'Movie'] | el][0] AS label\n",
    "LIMIT toInteger($limit)\n",
    "\"\"\"\n",
    "\n",
    "def get_candidates(input: str, type: str, limit: int = 3) -> List[Dict[str, str]]:\n",
    "    ft_query = generate_full_text_query(input)\n",
    "    candidates = graph.query(candidate_query, {'fulltextQuery': ft_query, 'index': type, 'limit':limit})\n",
    "    return candidates\n",
    "\n",
    "get_candidates(\"Matrex\", \"movie\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "VUs02iyCPEbW"
   },
   "outputs": [],
   "source": [
    "description_query = \"\"\"\n",
    "MATCH (m:Movie|Person)\n",
    "WHERE m.title = $candidate OR m.name = $candidate\n",
    "MATCH (m)-[r:ACTED_IN|DIRECTED|HAS_GENRE]-(t)\n",
    "WITH m, type(r) as type, collect(coalesce(t.name, t.title)) as names\n",
    "WITH m, type+\": \"+reduce(s=\"\", n IN names | s + n + \", \") as types\n",
    "WITH m, collect(types) as contexts\n",
    "WITH m, \"type:\" + labels(m)[0] + \"\\ntitle: \"+ coalesce(m.title, m.name) + \"\\nyear: \"+coalesce(m.released,\"\") +\"\\n\" +\n",
    "       reduce(s=\"\", c in contexts | s + substring(c, 0, size(c)-2) +\"\\n\") as context\n",
    "RETURN context LIMIT 1\n",
    "\"\"\"\n",
    "\n",
    "def get_information(entity: str, type:str) -> str:\n",
    "    candidates = get_candidates(entity, type)\n",
    "    if not candidates:\n",
    "        return \"No information was found about the movie or person in the database\"\n",
    "    elif len(candidates) > 1:\n",
    "        newline = '\\n'\n",
    "        return f\"Need additional information, which of these did you mean: {newline + newline.join(str(d) for d in candidates)}\"\n",
    "    data = graph.query(\n",
    "        description_query, params={\"candidate\": candidates[0]['candidate']}\n",
    "    )\n",
    "    return data[0][\"context\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "s0gaR6Hfx3F2",
    "outputId": "596ae85a-47ac-42c6-b016-b2fcdcdfd785"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Need additional information, which of these did you mean: \n",
      "{'candidate': 'John Lodge', 'label': 'Person'}\n",
      "{'candidate': 'John Warren', 'label': 'Person'}\n",
      "{'candidate': 'John Gray', 'label': 'Person'}\n"
     ]
    }
   ],
   "source": [
    "print(get_information(\"John\", \"person\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mIU44ESjOoIj",
    "outputId": "e175ba59-b8fa-460d-da0e-558b841aff20"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type:Actor\n",
      "title: Keanu Reeves\n",
      "year: \n",
      "ACTED_IN: Matrix Reloaded, The, Side by Side, Matrix Revolutions, The, Sweet November, Replacements, The, Hardball, Matrix, The, Constantine, Bill & Ted's Bogus Journey, Street Kings, Lake House, The, Chain Reaction, Walk in the Clouds, A, Little Buddha, Bill & Ted's Excellent Adventure, The Devil's Advocate, Johnny Mnemonic, Speed, Feeling Minnesota, The Neon Demon, 47 Ronin, Henry's Crime, Day the Earth Stood Still, The, John Wick, River's Edge, Man of Tai Chi, Dracula (Bram Stoker's Dracula), Point Break, My Own Private Idaho, Scanner Darkly, A, Something's Gotta Give, Watcher, The, Gift, The\n",
      "DIRECTED: Man of Tai Chi\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(get_information(\"Keanu Reeves\", \"person\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wZwinAkJPZ9Q",
    "outputId": "bef72b5a-f4e3-4774-9aa0-506b0c43b152"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No information was found about the movie or person in the database\n"
     ]
    }
   ],
   "source": [
    "print(get_information(\"Tomaz Bratanic\", \"person\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "arXgMCQ7PKoD"
   },
   "outputs": [],
   "source": [
    "recommendation_query_db_history = \"\"\"\n",
    "  MERGE (u:User {userId:$user_id})\n",
    "  WITH u\n",
    "  // get recommendation candidates\n",
    "  OPTIONAL MATCH (u)-[r1:RATED]->()<-[r2:RATED]-()-[r3:RATED]->(recommendation)\n",
    "  WHERE r1.rating > 3.5 AND r2.rating > 3.5 AND r3.rating > 3.5\n",
    "        AND NOT EXISTS {(u)-[:RATED]->(recommendation)}\n",
    "  // rank and limit recommendations\n",
    "  WITH u, recommendation, count(*) AS count\n",
    "  ORDER BY count DESC LIMIT 3\n",
    "  RETURN recommendation.title AS movie\n",
    "\"\"\"\n",
    "\n",
    "recommendation_query_genre = \"\"\"\n",
    "MATCH (m:Movie)-[:IN_GENRE]->(g:Genre {name:$genre})\n",
    "// filter out already seen movies by the user\n",
    "WHERE NOT EXISTS {\n",
    "  (m)<-[:RATED]-(:User {userId:$user_id})\n",
    "}\n",
    "// rank and limit recommendations\n",
    "WITH m\n",
    "ORDER BY m.imdbRating DESC LIMIT 3\n",
    "RETURN m.title AS movie\n",
    "\"\"\"\n",
    "\n",
    "def recommendation_query_movie(genre: bool) -> str:\n",
    "    return f\"\"\"\n",
    "MATCH (m1:Movie)<-[r1:RATED]-()-[r2:RATED]->(m2:Movie)\n",
    "WHERE r1.rating > 3.5 AND r2.rating > 3.5\n",
    "// filter out already seen movies by the user\n",
    "AND NOT EXISTS {{\n",
    "  (m2)<-[:RATED]-(:User {{userId:$user_id}})\n",
    "}}\n",
    "{'AND EXISTS {(m2)-[:IN_GENRE]->(:Genre {name:$genre})}' if genre else ''}\n",
    "// rank and limit recommendations\n",
    "WITH m2, count(*) AS count\n",
    "ORDER BY count DESC LIMIT 3\n",
    "RETURN m2.title As movie\n",
    "\"\"\"\n",
    "\n",
    "def get_user_id() -> int:\n",
    "    return 1\n",
    "\n",
    "def recommend_movie(movie: Optional[str] = None, genre: Optional[str] = None) -> str:\n",
    "    user_id = get_user_id()\n",
    "    params = {'user_id': user_id, 'genre':genre}\n",
    "    if not movie and not genre:\n",
    "        # Try to recommend a movie based on the information in the db\n",
    "        response = graph.query(recommendation_query_db_history, params)\n",
    "        try:\n",
    "            return \", \".join([el[\"movie\"] for el in response])\n",
    "        except:\n",
    "            return \"Can you tell us about some of the movies you liked?\"\n",
    "    if not movie and genre:\n",
    "        # Recommend top voted movies in the genre the user haven't seen before\n",
    "        response = graph.query(recommendation_query_genre, params)\n",
    "        try:\n",
    "            return \", \".join([el[\"movie\"] for el in response])\n",
    "        except:\n",
    "            return \"Something went wrong\"\n",
    "\n",
    "    candidates = get_candidates(movie, \"movie\")\n",
    "    if not candidates:\n",
    "        return \"The movie you mentioned wasn't found in the database\"\n",
    "    params['candidates'] = candidates\n",
    "    query = recommendation_query_movie(bool(genre))\n",
    "    response = graph.query(query, params)\n",
    "    try:\n",
    "        return \", \".join([el[\"movie\"] for el in response])\n",
    "    except:\n",
    "        return \"Something went wrong\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "NUOW_8ZCf7KU",
    "outputId": "43f86920-4e70-4567-94e8-9bc4c082f322"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'Star Wars: Episode IV - A New Hope, Star Wars: Episode V - The Empire Strikes Back, Star Wars: Episode VI - Return of the Jedi'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recommend_movie()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "o-ohHcbJhjD1",
    "outputId": "52135a59-f9cb-4f69-b4a8-529f352bc12b"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "\"Now You See Me 2, The Huntsman Winter's War, Kill Command\""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recommend_movie(genre=\"Action\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "Ik2p478n0Jo8",
    "outputId": "70a3226e-415e-45b6-e971-f065617c91fb"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'Pulp Fiction, Silence of the Lambs, The, Star Wars: Episode IV - A New Hope'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recommend_movie(movie=\"Top Gun\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "9ue8Oqwd0yAv",
    "outputId": "559c1a75-36a7-4b7a-d492-7cc7608aa58f"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'Star Wars: Episode IV - A New Hope, Star Wars: Episode V - The Empire Strikes Back, Matrix, The'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recommend_movie(movie=\"Top Gun\", genre=\"Action\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "2yKyzAlV1hkd"
   },
   "outputs": [],
   "source": [
    "store_rating_query = \"\"\"\n",
    "MERGE (u:User {userId:$user_id})\n",
    "WITH u\n",
    "UNWIND $candidates as candidate\n",
    "MATCH (m:Movie {title: candidate.candidate})\n",
    "MERGE (u)-[r:RATED]->(m)\n",
    "SET r.rating = 5.0\n",
    "RETURN distinct 'Noted' AS response\n",
    "\"\"\"\n",
    "def store_movie_rating(movie: str):\n",
    "    user_id = get_user_id()\n",
    "    candidates = get_candidates(movie, 'movie')\n",
    "    if not candidates:\n",
    "        return 'This movie is not in our database'\n",
    "    response = graph.query(store_rating_query, params={'user_id':user_id, 'candidates':candidates})\n",
    "    try:\n",
    "        return response[0]['response']\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return 'Something went wrong'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "VYYkr_kG3SjC",
    "outputId": "67fb0c84-657d-4a9a-ea80-a779ce596567"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'Noted'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "store_movie_rating('Top Gun')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "LqtrRwFt38HH"
   },
   "outputs": [],
   "source": [
    "from typing import Optional, Type\n",
    "\n",
    "from langchain.callbacks.manager import (\n",
    "    AsyncCallbackManagerForToolRun,\n",
    "    CallbackManagerForToolRun,\n",
    ")\n",
    "\n",
    "# Import things that are needed generically\n",
    "from langchain.pydantic_v1 import BaseModel, Field\n",
    "from langchain.tools import BaseTool, StructuredTool, tool\n",
    "\n",
    "class InformationInput(BaseModel):\n",
    "    entity: str = Field(description=\"movie or a person mentioned in the question\")\n",
    "    entity_type: str = Field(description=\"type of the entity. Available options are 'movie' or 'person'\")\n",
    "\n",
    "class InformationTool(BaseTool):\n",
    "    name = \"Information\"\n",
    "    description = \"useful for when you need to answer questions about various actors or movies\"\n",
    "    args_schema: Type[BaseModel] = InformationInput\n",
    "\n",
    "    def _run(\n",
    "        self, entity: str, entity_type: str, run_manager: Optional[CallbackManagerForToolRun] = None\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool.\"\"\"\n",
    "        return get_information(entity, entity_type)\n",
    "\n",
    "    async def _arun(\n",
    "        self,\n",
    "        entity: str,\n",
    "        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool asynchronously.\"\"\"\n",
    "        raise NotImplementedError(\"Calculator does not support async\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "qi-v-_jhQf5T"
   },
   "outputs": [],
   "source": [
    "all_genres = [\"Action\", \"Adventure\", \"Animation\", \"Children\", \"Comedy\", \"Crime\", \"Documentary\", \"Drama\", \"Fantasy\", \"Film-Noir\", \"Horror\", \"IMAX\", \"Musical\", \"Mystery\", \"Romance\", \"Sci-Fi\", \"Thriller\", \"War\", \"Western\"]\n",
    "\n",
    "\n",
    "class RecommenderInput(BaseModel):\n",
    "    movie: Optional[str] = Field(description=\"movie used for recommendation\")\n",
    "    genre: Optional[str] = Field(description=f\"genre used for recommendation. Available options are: {all_genres}\")\n",
    "\n",
    "\n",
    "class RecommenderTool(BaseTool):\n",
    "    name = \"Recommender\"\n",
    "    description = \"useful for when you need to recommend a movie\"\n",
    "    args_schema: Type[BaseModel] = RecommenderInput\n",
    "\n",
    "    def _run(\n",
    "        self, movie: Optional[str] = None, genre:Optional[str] = None, run_manager: Optional[CallbackManagerForToolRun] = None\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool.\"\"\"\n",
    "        return recommend_movie(movie, genre)\n",
    "\n",
    "    async def _arun(\n",
    "        self,\n",
    "        movie: Optional[str],\n",
    "        genre: Optional[str],\n",
    "        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool asynchronously.\"\"\"\n",
    "        raise NotImplementedError(\"Calculator does not support async\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "OS3cIpJwcC_e"
   },
   "outputs": [],
   "source": [
    "class MemoryInput(BaseModel):\n",
    "    movie: Optional[str] = Field(description=\"movie the user liked\")\n",
    "\n",
    "class MemoryTool(BaseTool):\n",
    "    name = \"Memory\"\n",
    "    description = \"useful for memorizing which movies the user liked\"\n",
    "    args_schema: Type[BaseModel] = MemoryInput\n",
    "\n",
    "    def _run(\n",
    "        self, movie: str, run_manager: Optional[CallbackManagerForToolRun] = None\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool.\"\"\"\n",
    "        return store_movie_rating(movie)\n",
    "\n",
    "    async def _arun(\n",
    "        self,\n",
    "        movie: str,\n",
    "        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,\n",
    "    ) -> str:\n",
    "        \"\"\"Use the tool asynchronously.\"\"\"\n",
    "        raise NotImplementedError(\"Calculator does not support async\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "WN1iLYGxPdP-"
   },
   "outputs": [],
   "source": [
    "from langchain.agents import Tool\n",
    "\n",
    "from langchain.tools.render import format_tool_to_openai_function\n",
    "\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(temperature=0, model=\"gpt-4\")\n",
    "tools = [\n",
    "    InformationTool(), RecommenderTool(), MemoryTool()\n",
    "]\n",
    "\n",
    "llm_with_tools = llm.bind(functions=[format_tool_to_openai_function(t) for t in tools])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "2kC7rkKGPhyt"
   },
   "outputs": [],
   "source": [
    "from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain.agents.format_scratchpad import format_to_openai_function_messages\n",
    "from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser\n",
    "from langchain.agents import AgentExecutor\n",
    "from langchain.schema import HumanMessage, AIMessage\n",
    "\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant that finds information about movies and recommends them. \"\n",
    "            \"If tools require follow up questions, make sure to ask the user for clarification. \"\n",
    "            \"Make sure to include any available options that need to be clarified in the follow up questions\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"user\", \"{input}\"),\n",
    "        MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "agent = (\n",
    "    {\n",
    "        \"input\": lambda x: x[\"input\"],\n",
    "        \"chat_history\": lambda x: x[\"chat_history\"] if x.get(\"chat_history\") else [],\n",
    "        \"agent_scratchpad\": lambda x: format_to_openai_function_messages(\n",
    "            x[\"intermediate_steps\"]\n",
    "        ),\n",
    "    }\n",
    "    | prompt\n",
    "    | llm_with_tools\n",
    "    | OpenAIFunctionsAgentOutputParser()\n",
    ")\n",
    "\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fYhfPtN-PkU_",
    "outputId": "cb5a44d2-d2e1-416f-f026-72880c692b3d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Information` with `{'entity': 'Pulp Fiction', 'entity_type': 'movie'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mtype:Movie\n",
      "title: Pulp Fiction\n",
      "year: 1994-10-14\n",
      "ACTED_IN: Tim Roth, John Travolta, Laura Lovelace, Samuel L. Jackson\n",
      "DIRECTED: Quentin Tarantino\n",
      "\u001b[0m\u001b[32;1m\u001b[1;3m\"Pulp Fiction\" was directed by Quentin Tarantino and the cast includes Tim Roth, John Travolta, Laura Lovelace, and Samuel L. Jackson.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'Who acted in pulp fiction?',\n",
       " 'output': '\"Pulp Fiction\" was directed by Quentin Tarantino and the cast includes Tim Roth, John Travolta, Laura Lovelace, and Samuel L. Jackson.'}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke({\"input\": \"Who acted in pulp fiction?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mC0zpnvRVPDv",
    "outputId": "8e502f97-d0bd-478a-da6e-c492085cf004"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Information` with `{'entity': 'John', 'entity_type': 'person'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mNeed additional information, which of these did you mean: \n",
      "{'candidate': 'John Lodge', 'label': 'Person'}\n",
      "{'candidate': 'John Warren', 'label': 'Person'}\n",
      "{'candidate': 'John Gray', 'label': 'Person'}\u001b[0m\u001b[32;1m\u001b[1;3mI found multiple people named John. Could you please specify which one you are interested in? Here are some options:\n",
      "\n",
      "1. John Lodge\n",
      "2. John Warren\n",
      "3. John Gray\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'What do you know about person John?',\n",
       " 'output': 'I found multiple people named John. Could you please specify which one you are interested in? Here are some options:\\n\\n1. John Lodge\\n2. John Warren\\n3. John Gray'}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke({\"input\": \"What do you know about person John?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hja3KiyVZYzp",
    "outputId": "b77e2ad0-b80b-423c-fb39-9db2f72ee15a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Information` with `{'entity': 'John Warren', 'entity_type': 'person'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mtype:Director\n",
      "title: John Warren\n",
      "year: \n",
      "DIRECTED: Major League: Back to the Minors\n",
      "\u001b[0m\u001b[32;1m\u001b[1;3mJohn Warren is a director known for his work on the movie \"Major League: Back to the Minors\".\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'John Warren',\n",
       " 'chat_history': [HumanMessage(content='What do you know about person John?'),\n",
       "  AIMessage(content='I found multiple people named John. Could you please specify which one you are interested in? Here are a few options:\\n\\n1. John Lodge\\n2. John Warren\\n3. John Gray')],\n",
       " 'output': 'John Warren is a director known for his work on the movie \"Major League: Back to the Minors\".'}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "history = [\n",
    "    HumanMessage(content=\"What do you know about person John?\"),\n",
    "    AIMessage(content=\"I found multiple people named John. Could you please specify which one you are interested in? Here are a few options:\\n\\n1. John Lodge\\n2. John Warren\\n3. John Gray\")\n",
    "]\n",
    "agent_executor.invoke({\"input\": \"John Warren\", \"chat_history\": history})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4HaasMGZPncM",
    "outputId": "da3af55f-286c-427d-faac-c37fee77316f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Recommender` with `{'genre': 'Comedy'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3mKeanu, Ice Age: The Great Egg-Scapade, Neighbors 2: Sorority Rising\u001b[0m\u001b[32;1m\u001b[1;3mSure, here are some comedy movies you might enjoy:\n",
      "\n",
      "1. Keanu\n",
      "2. Ice Age: The Great Egg-Scapade\n",
      "3. Neighbors 2: Sorority Rising\n",
      "\n",
      "Let me know if you need more recommendations!\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'Are there any good comedies out there?',\n",
       " 'output': 'Sure, here are some comedy movies you might enjoy:\\n\\n1. Keanu\\n2. Ice Age: The Great Egg-Scapade\\n3. Neighbors 2: Sorority Rising\\n\\nLet me know if you need more recommendations!'}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke(\n",
    "    {\"input\": \"Are there any good comedies out there?\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Vk34VR03aAVJ",
    "outputId": "d7950120-f9f6-4338-eca2-bb3e685bcd05"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Recommender` with `{'movie': 'Top Gun', 'genre': 'Action'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3mStar Wars: Episode IV - A New Hope, Star Wars: Episode V - The Empire Strikes Back, Matrix, The\u001b[0m\u001b[32;1m\u001b[1;3mBased on your interest in Top Gun, I would recommend the following action movies:\n",
      "\n",
      "1. Star Wars: Episode IV - A New Hope\n",
      "2. Star Wars: Episode V - The Empire Strikes Back\n",
      "3. The Matrix\n",
      "\n",
      "These movies have similar action-packed sequences and thrilling storylines. Enjoy your movie time!\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'I would like to see a good action movie like Top Gun?',\n",
       " 'output': 'Based on your interest in Top Gun, I would recommend the following action movies:\\n\\n1. Star Wars: Episode IV - A New Hope\\n2. Star Wars: Episode V - The Empire Strikes Back\\n3. The Matrix\\n\\nThese movies have similar action-packed sequences and thrilling storylines. Enjoy your movie time!'}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke(\n",
    "    {\"input\": \"I would like to see a good action movie like Top Gun?\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qxrAWIQcb37B",
    "outputId": "069d0c76-0aa7-41fe-dbfc-0b99a37e98da"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Memory` with `{'movie': 'raiders of the lost ark'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mNoted\u001b[0m\u001b[32;1m\u001b[1;3mThat's great! If you want, I can recommend similar movies to \"Raiders of the Lost Ark\". Would you like that?\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'I already watched raiders of the lost ark an I liked it.',\n",
       " 'output': 'That\\'s great! If you want, I can recommend similar movies to \"Raiders of the Lost Ark\". Would you like that?'}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke(\n",
    "    {\"input\": \"I already watched raiders of the lost ark and I liked it.\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KEZVzZOodCM1",
    "outputId": "79e5ec85-5de4-4254-e1e6-71f67caf80c4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Recommender` with `{'movie': 'Top Gun', 'genre': 'Action'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3mStar Wars: Episode IV - A New Hope, Star Wars: Episode V - The Empire Strikes Back, Matrix, The\u001b[0m\u001b[32;1m\u001b[1;3mBased on your interest in Top Gun, I would recommend the following action movies:\n",
      "\n",
      "1. Star Wars: Episode IV - A New Hope\n",
      "2. Star Wars: Episode V - The Empire Strikes Back\n",
      "3. The Matrix\n",
      "\n",
      "These are all highly rated and popular action movies. Enjoy your movie time!\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'I would like to see a good action movie like Top Gun?',\n",
       " 'output': 'Based on your interest in Top Gun, I would recommend the following action movies:\\n\\n1. Star Wars: Episode IV - A New Hope\\n2. Star Wars: Episode V - The Empire Strikes Back\\n3. The Matrix\\n\\nThese are all highly rated and popular action movies. Enjoy your movie time!'}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke(\n",
    "    {\"input\": \"I would like to see a good action movie like Top Gun?\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "FI7309DKdRpA"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyNPVQArh4o9RlIaws2LavM4",
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
